package KDTreeIndex;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.PriorityQueue;
import java.util.Queue;
import java.util.Set;
import java.util.TreeMap;

public class KDTree {

	private Node root;
	private Node current;
	private int dimension;
	private int deep;
	private int DataSize;
	private static final int MAX_SEARCH_STEP = 9;
	private HashMap<Key, ArrayList<Integer>> data;
	private HashMap<Integer, ArrayList<Node>> nodeList;
	private Queue<NodePriorityQueueElement> priorityQueue;
	private Queue<SearchPathQueueElement> searchPathQueue;

	/**
	 * 构造函数
	 * 
	 * @param Dimension
	 *            维数
	 */
	public KDTree(int Dimension) {
		dimension = Dimension;
		nodeList = new HashMap<Integer, ArrayList<Node>>();
	}

	/**
	 * 获取根节点
	 * 
	 * @return 根节点
	 */
	public Node getRoot() {
		return root;
	}

	/**
	 * 获取当前指向节点
	 * 
	 * @return 当前指向节点
	 */
	public Node getCurrent() {
		return current;
	}

	/**
	 * 构建KD-Tree
	 * 
	 * @param Data
	 */

	public ArrayList<Integer> getData(Key key) {
		return data.get(key);
	}

	public void buildTree(HashMap<Key, ArrayList<Integer>> Data) {
		data = Data;
		ArrayList<Key> keylist = convert(Data.keySet());

		root = new Node(dimension);
		root.setKeylist(keylist);
		current = root;
		creatTree(root, 0);
		deep = getMaxdeep(Data.size());
		DataSize = Data.size();
		for (int i = 0; i < deep; i++) {

			Iterator<Node> iterator = nodeList.get(i).iterator();

			while (iterator.hasNext()) {
				Node node = iterator.next();
				if (node.getKeylist().size() == 0) {
					continue;
				} else {
					creatTree(node, i + 1);
				}
			}
		}
	}

	/**
	 * 构造节点
	 * 
	 * @param parent
	 *            父节点
	 */
	private void creatTree(Node parent, int level) {
		if (parent.getKeylist().size() == 1) {
			if (parent.getKey().getKey(parent.getSplit()) > parent.getKeylist()
					.get(0).getKey(parent.getSplit())) {
				Node node = new Node(dimension);
				node.setParent(parent);
				node.setKey(parent.getKeylist().get(0));
				parent.setLeft(node);
			} else {
				Node node = new Node(dimension);
				node.setParent(parent);
				node.setKey(parent.getKeylist().get(0));
				parent.setRight(node);
			}
		} else {

			// 变量定义
			// splitList 分割列表
			TreeMap<Float, Integer> splitList = computeSplit(parent
					.getKeylist());

			parent.setSplitList(splitList);

			// 最大方差维
			int split = splitList.lastEntry().getValue();
			float splitKey = splitList.lastKey();
			// 左空间节点
			Node left = new Node(dimension);
			// 右空间节点
			Node right = new Node(dimension);
			// 左空间节点关键字列表
			ArrayList<Key> leftKeylist = new ArrayList<Key>();
			// 右空间节点关键字列表
			ArrayList<Key> rightKeylist = new ArrayList<Key>();
			// 父节点关键字列表 排序用

			// 设定左右空间划分参照维数
			left.setSplit(split);
			right.setSplit(split);
			// 父节点关联
			left.setParent(parent);
			right.setParent(parent);
			// 左右空间关键字列表关联

			parent.setLeft(left);
			parent.setRight(right);

			left.setKeylist(leftKeylist);
			right.setKeylist(rightKeylist);

			Key centerKey = getCenterkey(parent.getKeylist(), splitList);
			parent.setKey(centerKey);
			parent.getKeylist().remove(centerKey);
			// 划分左右空间

			for (Key key : parent.getKeylist()) {
				if (key.getKey(split) > centerKey.getKey(split)) {
					rightKeylist.add(key);
					continue;
				} else if (key.getKey(split) < centerKey.getKey(split)) {
					leftKeylist.add(key);

					continue;
				} else {

					// 垂直于分割超平面的方向轴序号上的关键字数值相同 则转向 方差较小的另一维
					boolean flag = true;
					while (flag) {
						int splitTemp = splitList.lowerEntry(splitKey)
								.getValue();
						splitKey = splitList.lowerKey(splitKey);
						if (key.getKey(splitTemp) > centerKey.getKey(splitTemp)) {
							rightKeylist.add(key);
							flag = false;
							splitKey = splitList.lastKey();

							continue;
						} else if (key.getKey(splitTemp) < centerKey
								.getKey(splitTemp)) {
							leftKeylist.add(key);
							flag = false;
							splitKey = splitList.lastKey();

							continue;
						} else {
						}
					}

				}
				// 父点击完成左右空间划分后 清除关键字列表 节省空间

			}
			parent.clearKeylist();
			nodeList.put(level, new ArrayList<Node>());

			nodeList.get(level).add(left);
			nodeList.get(level).add(right);
			if (leftKeylist.size() == 0 && rightKeylist.size() == 1) {

				right.setKey(rightKeylist.get(0));

			} else if (rightKeylist.size() == 0 && leftKeylist.size() == 1) {

				left.setKey(leftKeylist.get(0));
			}
		}

	}

	/**
	 * 计算Split （垂直于分割超平面的方向轴序号）
	 * 
	 * @param keylist
	 *            关键字列表
	 * @return 分割维数TreeMap（方差排序列表）
	 */
	private TreeMap<Float, Integer> computeSplit(ArrayList<Key> keylist) {
		TreeMap<Float, Integer> splitList = new TreeMap<Float, Integer>();
		float[] avg = computeAvg(keylist);
		float[] var = new float[dimension];
		for (int i = 0; i < keylist.size(); i++) {
			for (int j = 0; j < dimension; j++) {
				float temp = keylist.get(i).getKey(j) - avg[j];
				var[j] += temp * temp;
			}
		}
		for (int i = 0; i < var.length; i++) {
			splitList.put(var[i], i);
		}
		return splitList;
	}

	/**
	 * 计算平均数
	 * 
	 * @param keylist
	 *            关键字列表
	 * @return 各个维数上的平均数
	 */
	private float[] computeAvg(ArrayList<Key> keylist) {
		float[] sum = new float[dimension];
		float[] avg = new float[dimension];
		for (int i = 0; i < keylist.size(); i++) {
			for (int j = 0; j < dimension; j++) {
				sum[j] += keylist.get(i).getKey(j);
			}
		}
		for (int j = 0; j < dimension; j++) {
			avg[j] = sum[j] / keylist.size();
		}
		return avg;
	}

	// ArrayList<Key> 转 Set<Key>
	/**
	 * ArrayList<Key> 转 Set<Key>
	 * 
	 * @param keyset
	 *            ArrayList<Key>
	 * @return Set<Key>
	 */
	private ArrayList<Key> convert(Set<Key> keyset) {
		ArrayList<Key> keylist = new ArrayList<Key>();
		Iterator<Key> iterator = keyset.iterator();
		while (iterator.hasNext()) {
			keylist.add(iterator.next());
		}
		return keylist;
	}

	private int getMaxdeep(int N) {

		double n = Math.log(N) / Math.log(2);
		return (int) n;
	}

	private Key getCenterkey(ArrayList<Key> keylist,
			TreeMap<Float, Integer> splitList) {
		TreeMap<Float, Key> temp = new TreeMap<Float, Key>();
		if (keylist.size() == 1) {
			return keylist.get(0);
		} else {
			for (Key key : keylist) {
				float num = 0;
				for (int i : splitList.values()) {
					num += key.getKey(i);
					num *= 0.1;
				}
				num = num * 10;
				temp.put(num, key);
			}
			float lastkey = temp.lastKey();
			for (int i = 0; i < temp.size() / 2; i++) {
				lastkey = temp.lowerEntry(lastkey).getKey();
			}
			return temp.get(lastkey);
		}
	}

	private Node goToLeft(Node node, Key target, int Split) {
		NodePriorityQueueElement queueNode = new NodePriorityQueueElement();
		SearchPathQueueElement searchNode = new SearchPathQueueElement();

		queueNode.setNode(node.getRight());
		searchNode.setNode(node.getLeft(), target);

		int eigenvalue = Math.abs(target.getKey(Split)
				- node.getRight().getKey().getKey(Split));
		queueNode.setEigenvalue(eigenvalue);

		priorityQueue.offer(queueNode);
		searchPathQueue.offer(searchNode);

		return node.getLeft();
	}

	private Node goToRight(Node node, Key target, int Split) {
		NodePriorityQueueElement queueNode = new NodePriorityQueueElement();
		SearchPathQueueElement searchNode = new SearchPathQueueElement();

		queueNode.setNode(node.getLeft());
		searchNode.setNode(node.getRight(), target);

		int eigenvalue = Math.abs(target.getKey(Split)
				- node.getRight().getKey().getKey(Split));
		queueNode.setEigenvalue(eigenvalue);

		priorityQueue.offer(queueNode);
		searchPathQueue.offer(searchNode);

		return node.getRight();
	}

	private Node exploreToLeaf(Node node, Key target) {
		while (node.getLeft() != null && node.getRight() != null) {
			int Split = node.getSplit();
			TreeMap<Float, Integer> splitList = node.getSplitList();
			float splitKey = splitList.lastKey();
			if(node.getKey().equals(target))
			{
				return node;
			}
			if (target.getKey(Split) > node.getKey().getKey(Split)) {

				node = goToRight(node, target, Split);
			} else if (target.getKey(Split) < this.current.getKey().getKey(
					Split)) {
				node = goToLeft(node, target, Split);
			} else {
				boolean flag = true;
				while (flag) {

					int splitTemp = splitList.lowerEntry(splitKey).getValue();
					splitKey = splitList.lowerKey(splitKey);
					if (target.getKey(splitTemp) > this.current.getKey()
							.getKey(splitTemp)) {

						flag = false;
						node = goToRight(node, target, Split);
						continue;
					} else if (target.getKey(splitTemp) < this.current.getKey()
							.getKey(splitTemp)) {

						flag = false;
						node = goToLeft(node, target, Split);
						continue;
					}
				}
			}
		}
		return node;
	}

	public ArrayList<Integer> Search(Key target) {
		if (this.root == null) {
			return new ArrayList<Integer>();
		} else {
			int step = 0;

			Comparator<NodePriorityQueueElement> nodePriorityQueueOder = new Comparator<NodePriorityQueueElement>() {

				@Override
				public int compare(NodePriorityQueueElement o1,
						NodePriorityQueueElement o2) {
					if (o1.getEigenvalue() < o2.getEigenvalue()) {
						return -1;
					} else if (o1.getEigenvalue() > o2.getEigenvalue()) {
						return 1;
					} else {
						return 0;
					}
				}
			};

			Comparator<SearchPathQueueElement> searchPathOder = new Comparator<SearchPathQueueElement>() {

				@Override
				public int compare(SearchPathQueueElement o1,
						SearchPathQueueElement o2) {
					if (o1.getDistance() < o2.getDistance()) {
						return -1;
					} else if (o1.getDistance() > o2.getDistance()) {
						return 1;
					} else {
						return 0;
					}
				}
			};
			priorityQueue = new PriorityQueue<NodePriorityQueueElement>(
					deep + 1, nodePriorityQueueOder);
			searchPathQueue = new PriorityQueue<SearchPathQueueElement>(
					DataSize, searchPathOder);

			SearchPathQueueElement searchNode = new SearchPathQueueElement();
			NodePriorityQueueElement queueRootNode = new NodePriorityQueueElement();

			queueRootNode.setNode(root);
			searchNode.setNode(root, target);

			searchPathQueue.offer(searchNode);
			priorityQueue.offer(queueRootNode);

			while (priorityQueue.size() > 0 && step < MAX_SEARCH_STEP) {
				exploreToLeaf(priorityQueue.remove().getNode(), target);
			}
		}
		return this.data.get(searchPathQueue.remove().getNode().getKey());

	}

}

class NodePriorityQueueElement implements Comparable<NodePriorityQueueElement> {
	/**
		 * 
		 */
	private int eigenvalue;
	private Node node;

	protected int getEigenvalue() {
		return eigenvalue;
	}

	protected Node getNode() {
		return node;
	}

	protected void setEigenvalue(int eigenvalue) {
		this.eigenvalue = eigenvalue;
	}

	protected void setNode(Node node) {
		this.node = node;
	}

	@Override
	public String toString() {
		return node.getKey().toString() + "," + eigenvalue;
	}

	@Override
	public int compareTo(NodePriorityQueueElement arg0) {
		if (arg0.eigenvalue < this.eigenvalue) {
			return -1;
		} else if (arg0.eigenvalue > this.eigenvalue) {
			return 1;
		} else {
			return 0;
		}
	}
}

class SearchPathQueueElement implements Comparable<SearchPathQueueElement> {
	private double distance;
	private Node node;

	protected double getDistance() {
		return distance;
	}

	protected Node getNode() {
		return node;
	}

	protected void setDistance(double distance) {
		this.distance = distance;
	}

	protected void setNode(Node node, Key Targer) {
		this.node = node;
		distance = computeDistance(node.getKey(), Targer);
	}

	@Override
	public String toString() {
		return node.getKey().toString() + "," + distance;
	}

	private double computeDistance(Key key1, Key key2) {
		double distance = 0.0f;
		for (int i = 0; i < key1.getKey().length; i++) {
			float temp = key1.getKey(i) - key2.getKey(i);
			distance += temp * temp;
		}
		return Math.sqrt(distance);
	}

	@Override
	public int compareTo(SearchPathQueueElement o) {
		if (o.distance > this.distance) {
			return 1;
		} else if (o.distance < this.distance) {
			return -1;
		} else {
			return 0;
		}
	}
}
